from .utils.register import Register
from .ops import *
import logging
import torch

current_work_registrar = None
logger = logging.getLogger(__name__)

def enable(
    lib=None,
    unused=None,
    registrar=None,
):
    if registrar is None:
        logger.warning("registrar is none")
        return

    global current_work_registrar
    current_work_registrar = registrar(
        (
            ("_softmax", softmax_new),
        ),
        user_unused_ops_list = [] if unused is None else unused,
        lib = lib,
    )

class use_custom_ops:
    def __init__(self, unused=None):
        self.lib = torch.library.Library("aten", "IMPL")
        self.unused = [] if unused is None else unused
        self.registrar = Register

    def __enter__(self):
        enable(
            lib=self.lib,
            unused=self.unused,
            registrar=self.registrar,
        )

    def __exit__(self, exc_type, exc_val, exc_tb):
        global current_work_registrar
        del self.lib
        del self.unused
        del self.registrar
        del current_work_registrar

def test_compiling():
    print("compile triton kernels with triton on ascend compiler, use MLIR_ENABLE_DUMP=1 to see the more details.\n")

    for key, test_func in compiled_kernel_list.items():
        print(f"compile {key} ...")
        test_func()

__all__ = [
    "enable",
    "use_custom_ops",
    "test_compiling"
]